{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2021-2022 @ Shenzhen Bay Laboratory & Peking University & Huawei Technologies Co., Ltd\n",
    "\n",
    "This code is a part of Cybertron package.\n",
    "\n",
    "The Cybertron is open-source software based on the AI-framework:\n",
    "MindSpore (https://www.mindspore.cn/)\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "\n",
    "You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n",
    "\n",
    "Cybertron tutorial 07: Dataset with force"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(16568:139686452430656,MainProcess):2022-08-10-17:35:37.767.789 [mindspore/run_check/_check_version.py:137] Can not found cuda libs, please confirm that the correct cuda version has been installed, you can refer to the installation guidelines: https://www.mindspore.cn/install\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import mindspore as ms\n",
    "from mindspore import nn\n",
    "from mindspore import Tensor\n",
    "from mindspore import dataset as ds\n",
    "from mindspore import context\n",
    "from mindspore.train import Model\n",
    "from mindspore.train.callback import ModelCheckpoint, CheckpointConfig\n",
    "\n",
    "from cybertron import Cybertron\n",
    "from cybertron import MolCT\n",
    "from cybertron import AtomwiseReadout\n",
    "from cybertron.train import MAE, RMSE, MLoss, MSELoss\n",
    "from cybertron.train import WithForceLossCell, WithForceEvalCell\n",
    "from cybertron.train import TrainMonitor\n",
    "from cybertron.train import TransformerLR\n",
    "\n",
    "context.set_context(mode=context.GRAPH_MODE, device_target=\"GPU\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys_name = 'dataset_ethanol_normed_'\n",
    "\n",
    "train_file = sys_name + 'trainset_1024.npz'\n",
    "valid_file = sys_name + 'validset_128.npz'\n",
    "\n",
    "train_data = np.load(train_file)\n",
    "valid_data = np.load(valid_file)\n",
    "\n",
    "atom_types = Tensor(train_data['Z'], ms.int32)\n",
    "scale = train_data['scale']\n",
    "shift = train_data['shift']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "mod = MolCT(\n",
    "    cutoff=1,\n",
    "    n_interaction=3,\n",
    "    dim_feature=128,\n",
    "    n_heads=8,\n",
    "    max_cycles=1,\n",
    "    fixed_cycles=True,\n",
    "    length_unit='nm',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "readout = AtomwiseReadout(mod, dim_output=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] PIPELINE(16568,7f0b49608740,python):2022-08-10-17:35:44.933.029 [mindspore/ccsrc/pipeline/jit/pipeline.cc:169] CheckArgValid] The data types of Tensor:[[ True  True  True  True  True  True  True  True  True]] is bool, which may cause SelectKernelInfo failure for operator [AddN]. For more details, please refer to the FAQ at https://www.mindspore.cn.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================================================================\n",
      "Cybertron Engine, Ride-on!\n",
      "--------------------------------------------------------------------------------\n",
      "    Length unit: nm\n",
      "    Input unit scale: 1\n",
      "    Using fixed atom type index:\n",
      "       Atom 0:     6\n",
      "       Atom 1:     6\n",
      "       Atom 2:     8\n",
      "       Atom 3:     1\n",
      "       Atom 4:     1\n",
      "       Atom 5:     1\n",
      "       Atom 6:     1\n",
      "       Atom 7:     1\n",
      "       Atom 8:     1\n",
      "--------------------------------------------------------------------------------\n",
      "    Deep molecular model:  MolCT\n",
      "--------------------------------------------------------------------------------\n",
      "       Length unit: nm\n",
      "       Atom embedding size: 64\n",
      "       Cutoff distance: 1.0 nm\n",
      "       Radical basis function (RBF): LogGaussianBasis\n",
      "          Minimum distance: 0.04 nm\n",
      "          Maximum distance: 1.0 nm\n",
      "          Reference distance: 1.0 nm\n",
      "          Log Gaussian begin: -3.218876\n",
      "          Log Gaussian end: 0.006724119\n",
      "          Interval for log Gaussian: 0.0512\n",
      "          Sigma for log gaussian: 0.3\n",
      "          Number of basis functions: 64\n",
      "          Rescale the range of RBF to (-1,1).\n",
      "       Calculate distance: Yes\n",
      "       Calculate bond: No\n",
      "       Feature dimension: 128\n",
      "--------------------------------------------------------------------------------\n",
      "       Using 3 independent interaction layers:\n",
      "--------------------------------------------------------------------------------\n",
      "       0. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "       1. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "       2. Neural Interaction Unit\n",
      "          Feature dimension: 128\n",
      "          Activation function: Swish\n",
      "          Encoding distance: Yes\n",
      "          Encoding bond: No\n",
      "          Number of heads in multi-haed attention: 8\n",
      "          Use feed forward network: No\n",
      "--------------------------------------------------------------------------------\n",
      "    Readout network: AtomwiseReadout\n",
      "--------------------------------------------------------------------------------\n",
      "       Activation function: Swish\n",
      "       Decoder: HalveDecoder\n",
      "       Aggregator: TensorSummation\n",
      "       Representation dimension: 128\n",
      "       Readout dimension: 1\n",
      "       Scale: 1.0\n",
      "       Shift: 0.0\n",
      "       Scaleshift mode: Atomwise\n",
      "       Reference value for atom types: None\n",
      "       Output unit: None\n",
      "       Reduce axis: -2\n",
      "--------------------------------------------------------------------------------\n",
      "    Output dimension: 1\n",
      "    Output unit for Cybertron: None\n",
      "    Output unit scale: 1.0\n",
      "================================================================================\n"
     ]
    }
   ],
   "source": [
    "net = Cybertron(mod, readout=readout, atom_types=atom_types, length_unit='nm')\n",
    "net.print_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 model.atom_embedding.embedding_table (64, 128)\n",
      "1 model.dis_filter.linear.weight (128, 64)\n",
      "2 model.dis_filter.linear.bias (128,)\n",
      "3 model.dis_filter.residual.nonlinear.mlp.0.weight (128, 128)\n",
      "4 model.dis_filter.residual.nonlinear.mlp.0.bias (128,)\n",
      "5 model.dis_filter.residual.nonlinear.mlp.1.weight (128, 128)\n",
      "6 model.dis_filter.residual.nonlinear.mlp.1.bias (128,)\n",
      "7 model.interactions.0.positional_embedding.norm.gamma (128,)\n",
      "8 model.interactions.0.positional_embedding.norm.beta (128,)\n",
      "9 model.interactions.0.positional_embedding.x2q.weight (128, 128)\n",
      "10 model.interactions.0.positional_embedding.x2k.weight (128, 128)\n",
      "11 model.interactions.0.positional_embedding.x2v.weight (128, 128)\n",
      "12 model.interactions.0.multi_head_attention.output.weight (128, 128)\n",
      "13 model.interactions.1.positional_embedding.norm.gamma (128,)\n",
      "14 model.interactions.1.positional_embedding.norm.beta (128,)\n",
      "15 model.interactions.1.positional_embedding.x2q.weight (128, 128)\n",
      "16 model.interactions.1.positional_embedding.x2k.weight (128, 128)\n",
      "17 model.interactions.1.positional_embedding.x2v.weight (128, 128)\n",
      "18 model.interactions.1.multi_head_attention.output.weight (128, 128)\n",
      "19 model.interactions.2.positional_embedding.norm.gamma (128,)\n",
      "20 model.interactions.2.positional_embedding.norm.beta (128,)\n",
      "21 model.interactions.2.positional_embedding.x2q.weight (128, 128)\n",
      "22 model.interactions.2.positional_embedding.x2k.weight (128, 128)\n",
      "23 model.interactions.2.positional_embedding.x2v.weight (128, 128)\n",
      "24 model.interactions.2.multi_head_attention.output.weight (128, 128)\n",
      "25 readout.decoder.output.mlp.0.weight (64, 128)\n",
      "26 readout.decoder.output.mlp.0.bias (64,)\n",
      "27 readout.decoder.output.mlp.1.weight (1, 64)\n",
      "28 readout.decoder.output.mlp.1.bias (1,)\n",
      "Total parameters:  255233\n"
     ]
    }
   ],
   "source": [
    "tot_params = 0\n",
    "for i, param in enumerate(net.trainable_params()):\n",
    "    tot_params += param.size\n",
    "    print(i, param.name, param.shape)\n",
    "print('Total parameters: ', tot_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epoch = 8\n",
    "repeat_time = 1\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_train = ds.NumpySlicesDataset(\n",
    "    {'R': train_data['R'], 'F': train_data['F'], 'E': train_data['E']}, shuffle=True)\n",
    "ds_train = ds_train.batch(batch_size)\n",
    "ds_train = ds_train.repeat(repeat_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_valid = ds.NumpySlicesDataset(\n",
    "    {'R': valid_data['R'], 'F': valid_data['F'], 'E': valid_data['E']}, shuffle=False)\n",
    "ds_valid = ds_valid.batch(128)\n",
    "ds_valid = ds_valid.repeat(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WithForceLossCell with input type: RFE\n",
      "WithForceEvalCell with input type: RFE\n",
      "   with scaleshift for training and evaluate dataset:\n",
      "   Scale: [1.9488194]\n",
      "   Shift: [-45185.31]\n",
      "   Scaleshift mode: atomwise\n"
     ]
    }
   ],
   "source": [
    "force_dis = train_data['avg_force_dis']\n",
    "loss_fn = MSELoss(ratio_energy=1, ratio_forces=100, force_dis=force_dis)\n",
    "loss_network = WithForceLossCell('RFE', net, loss_fn)\n",
    "eval_network = WithForceEvalCell('RFE', net, loss_fn, scale=scale, shift=shift)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = TransformerLR(learning_rate=1., warmup_steps=4000, dimension=128)\n",
    "optim = nn.Adam(params=net.trainable_params(), learning_rate=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_mae = 'EnergyMAE'\n",
    "forces_mae = 'ForcesMAE'\n",
    "forces_rmse = 'ForcesRMSE'\n",
    "eval_loss = 'EvalLoss'\n",
    "model = Model(loss_network, eval_network=eval_network, optimizer=optim,\n",
    "              metrics={eval_loss: MLoss(), energy_mae: MAE([1, 2]), forces_mae: MAE([3, 4]),\n",
    "                       forces_rmse: RMSE([3, 4], atom_aggregate='sum')})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "outdir = 'Tutorial_C07'\n",
    "outname = outdir + '_' + net.model_name\n",
    "record_cb = TrainMonitor(model, outname, per_epoch=1, avg_steps=32,\n",
    "                         directory=outdir, eval_dataset=ds_valid, best_ckpt_metrics=forces_rmse)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_ck = CheckpointConfig(save_checkpoint_steps=32, keep_checkpoint_max=64, append_info=[net.hyper_param])\n",
    "ckpoint_cb = ModelCheckpoint(prefix=outname, directory=outdir, config=config_ck)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[WARNING] ME(16568:139686452430656,MainProcess):2022-08-10-17:35:55.296.076 [mindspore/train/model.py:1097] For TrainMonitor callback, {'epoch_end', 'step_end', 'begin'} methods may not be supported in later version, Use methods prefixed with 'on_train' or 'on_eval' instead when using customized callbacks.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start training ...\n",
      "Epoch: 1, Step: 32, Learning_rate: 1.0830951e-05, Last_Loss: 1.2773383, Avg_loss: 1.3382577523589134, EvalLoss: 1.350154161453247, EnergyMAE: 17.78466796875, ForcesMAE: 816.9448784722222, ForcesRMSE: 1895.1583926768058\n",
      "Epoch: 2, Step: 64, Learning_rate: 2.2011289e-05, Last_Loss: 1.0896915, Avg_loss: 1.297719169408083, EvalLoss: 1.3006870746612549, EnergyMAE: 13.238037109375, ForcesMAE: 802.8101671006945, ForcesRMSE: 1863.905457783617\n",
      "Epoch: 3, Step: 96, Learning_rate: 3.3191627e-05, Last_Loss: 0.8397354, Avg_loss: 1.1701170392334461, EvalLoss: 0.9390307664871216, EnergyMAE: 13.350341796875, ForcesMAE: 670.5280490451389, ForcesRMSE: 1581.9937491097182\n",
      "Epoch: 4, Step: 128, Learning_rate: 4.437197e-05, Last_Loss: 0.5989347, Avg_loss: 0.689078314229846, EvalLoss: 0.5663008689880371, EnergyMAE: 10.6328125, ForcesMAE: 504.49359809027777, ForcesRMSE: 1227.509764613798\n",
      "Epoch: 5, Step: 160, Learning_rate: 5.5552304e-05, Last_Loss: 0.3392536, Avg_loss: 0.45017723087221384, EvalLoss: 0.3031836748123169, EnergyMAE: 10.0458984375, ForcesMAE: 375.38783094618054, ForcesRMSE: 895.3201165567046\n",
      "Epoch: 6, Step: 192, Learning_rate: 6.6732646e-05, Last_Loss: 0.15741858, Avg_loss: 0.22514202212914824, EvalLoss: 0.1697137951850891, EnergyMAE: 12.449462890625, ForcesMAE: 278.5558810763889, ForcesRMSE: 659.8750302561505\n",
      "Epoch: 7, Step: 224, Learning_rate: 7.791298e-05, Last_Loss: 0.10840593, Avg_loss: 0.14769150223582983, EvalLoss: 0.11459552496671677, EnergyMAE: 7.0234375, ForcesMAE: 232.14268663194446, ForcesRMSE: 548.4206617389814\n",
      "Epoch: 8, Step: 256, Learning_rate: 8.9093315e-05, Last_Loss: 0.07344754, Avg_loss: 0.10658725094981492, EvalLoss: 0.10072694718837738, EnergyMAE: 14.091796875, ForcesMAE: 213.02398003472223, ForcesRMSE: 497.3395609529659\n",
      "Training Fininshed!\n",
      "Training Time: 00:00:24\n"
     ]
    }
   ],
   "source": [
    "print(\"Start training ...\")\n",
    "beg_time = time.time()\n",
    "model.train(n_epoch, ds_train, callbacks=[\n",
    "    record_cb, ckpoint_cb], dataset_sink_mode=False)\n",
    "end_time = time.time()\n",
    "used_time = end_time - beg_time\n",
    "m, s = divmod(used_time, 60)\n",
    "h, m = divmod(m, 60)\n",
    "print(\"Training Fininshed!\")\n",
    "print(\"Training Time: %02d:%02d:%02d\" % (h, m, s))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.5 ('mindsponge')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2496ecc683137a232cae2452fbbdd53dab340598b6e499c8995be760f3a431b4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
